#include <Arduino.h>
#include <WiFi.h>
#include "atom.h"
#include "wlan.h"
#include "syslog.h"
#include "console.h"
#include "esp_wifi.h"
#include "drv_wlan.h"
#include "debug.h"
#include <ESP32Ping.h>  // 使用ESP32Ping库
#include "lwip/inet.h"
#include "lwip/err.h"
#include "lwip/sys.h"
#include "e2p.h"
#include <esp_system.h>
#include <esp_crc.h>  // Add this for esp_crc32_le

static WlanConfig _config = {
    NULL, 0, false
};

static Wlan _wlan;

// 声明外部E2P设备
static E2prom *g_e2p = NULL;  // 假设E2P设备已经在其他地方创建

// 计算CRC32
static uint32_t calculate_crc32(const void* data, size_t length) {
    return esp_crc32_le(0, (const uint8_t*)data, length);
}

// 实现配置读取
static int esp32_read_config(Wlan *dev, char *ssid, char *password) {
    WlanStorageConfig config;
    E2pDataBlock data = {
        .addr = WLAN_CONFIG_ADDR,
        .buffer = (uint8_t*)&config,
        .size = sizeof(config)
    };
    
    // 从E2P读取配置
    if (E2PRead(g_e2p, &data) != OK) {
        return NG;
    }
    
    // 验证CRC
    uint32_t calc_crc = calculate_crc32(&config, sizeof(config) - sizeof(config.crc));
    if (calc_crc != config.crc) {
        return NG;  // CRC校验失败
    }
    
    // 复制配置
    strncpy(ssid, config.ssid, 32);
    strncpy(password, config.password, 64);
    
    return OK;
}

// 实现配置保存
static int esp32_save_config(Wlan *dev, const char *ssid, const char *password) {
    WlanStorageConfig config;
    
    // 清空并填充配置
    memset(&config, 0, sizeof(config));
    strncpy(config.ssid, ssid, 32);
    strncpy(config.password, password, 64);
    
    // 计算CRC
    config.crc = calculate_crc32(&config, sizeof(config) - sizeof(config.crc));
    
    // 准备数据块
    E2pDataBlock data = {
        .addr = WLAN_CONFIG_ADDR,
        .buffer = (uint8_t*)&config,
        .size = sizeof(config)
    };
    
    // 写入E2P并立即提交
    return E2PWrite(g_e2p, &data, 1);
}

// 实现MAC地址获取
static int esp32_get_mac(Wlan *dev, uint8_t *mac) {
    if (!mac) return NG;
    
    // 获取ESP32的MAC地址
    esp_read_mac(mac, ESP_MAC_WIFI_STA);
    return OK;
}

// WLAN operations implementation
static int esp32_set_mode(Wlan *dev, WlanMode_t mode)
{
    switch(mode) {
        case WLAN_MODE_NONE:
            WiFi.mode(WIFI_OFF);
            break;
        case WLAN_MODE_STATION:
            WiFi.mode(WIFI_STA);
            break;
        case WLAN_MODE_AP:
            WiFi.mode(WIFI_AP);
            break;
        default:
            return NG;
    }
    
    dev->Protected.mode = mode;
    return OK;
}

static WlanMode_t esp32_get_mode(Wlan *dev)
{
    wifi_mode_t mode;
    if(esp_wifi_get_mode(&mode) != ESP_OK) {
        return WLAN_MODE_NONE;
    }

    switch(mode) {
        case WIFI_MODE_NULL:
            return WLAN_MODE_NONE;
        case WIFI_MODE_STA:
            return WLAN_MODE_STATION;
        case WIFI_MODE_AP:
            return WLAN_MODE_AP;
        default:
            return WLAN_MODE_NONE;
    }
}

static int esp32_scan(Wlan *dev, WlanScanInfo *info, int max_results)
{
    WlanConfig* config = (WlanConfig*)dev->Protected.config;
    
    if(dev->Protected.mode != WLAN_MODE_STATION) {
        SYSLOG(LOG_ERR, "WLAN must be in station mode for scanning");
        return NG;
    }

    if(config->is_scanning) {
        SYSLOG(LOG_ERR, "Scan already in progress");
        return NG;
    }

    config->is_scanning = true;
    int n = WiFi.scanNetworks();
    config->is_scanning = false;

    if(n < 0) {
        SYSLOG(LOG_ERR, "Scan failed");
        return NG;
    }

    // Store scan results
    for(int i = 0; i < n && i < max_results; i++) {
        // Copy SSID
        strncpy((char*)info[i].ssid, WiFi.SSID(i).c_str(), 32);
        info[i].ssid[32] = '\0';
        
        // Copy BSSID (MAC address)
        uint8_t* bssid = WiFi.BSSID(i);
        memcpy(info[i].bssid, bssid, 6);
        
        // Copy channel
        info[i].channel = WiFi.channel(i);
    }

    return n;
}

static int esp32_connect(Wlan *dev, const char *ssid, const char *password, uint32_t timeout_sec)
{
    if(dev->Protected.mode != WLAN_MODE_STATION) {
        SYSLOG(LOG_ERR, "WLAN must be in station mode for connecting");
        return NG;
    }

    WiFi.begin(ssid, password);
    
    // Wait for connection with specified timeout
    uint32_t timeout_count = timeout_sec; // 每100ms检查一次，所以超时次数需要乘以10
    while(WiFi.status() != WL_CONNECTED && timeout_count > 0) {
        atomTimerDelay(1000);
        timeout_count--;
    }

    if(WiFi.status() != WL_CONNECTED) {
        SYSLOG(LOG_ERR, "Failed to connect to WiFi: timeout after %d seconds", timeout_sec);
        return NG;
    }

    return OK;
}

static int esp32_startap(Wlan *dev, const char *ssid, const char *password)
{
    if(dev->Protected.mode != WLAN_MODE_AP) {
        SYSLOG(LOG_ERR, "WLAN must be in AP mode for starting AP");
        return NG;
    }

    IPAddress ap_ip, subnet, gateway;
    ap_ip.fromString(BSP_USE_WIFI_AP_IP);
    subnet.fromString(BSP_USE_WIFI_AP_SUBNET);
    gateway.fromString(BSP_USE_WIFI_AP_GATEWAY);


    if(!WiFi.softAPConfig(ap_ip, gateway, subnet)) {
        SYSLOG(LOG_ERR, "AP Config failed");
        return NG;
    }

    if(!WiFi.softAP(ssid, password)) {
        SYSLOG(LOG_ERR, "AP Start failed");
        return NG;
    }

    return OK;
}

static int esp32_disconnect(Wlan *dev)
{
    if(dev->Protected.mode == WLAN_MODE_STATION) {
        WiFi.disconnect();
    } else if(dev->Protected.mode == WLAN_MODE_AP) {
        WiFi.softAPdisconnect();
    }
    return OK;
}

// 添加辅助函数用于IP地址拷贝
static void copy_ip_address(uint8_t* dest, const IPAddress& src) {
    for(int i = 0; i < 4; i++) {
        dest[i] = src[i];
    }
}

static int esp32_get_net_info(Wlan *dev, WlanNetInfo *info)
{
    if (!info) return NG;
    
    // 清零结构体
    memset(info, 0, sizeof(WlanNetInfo));
    
    WlanMode_t mode = esp32_get_mode(dev);
    if (mode == WLAN_MODE_STATION) {
        // Station 模式
        if (WiFi.status() == WL_CONNECTED) {
            info->connected = 1;
            
            // 使用新的辅助函数正确拷贝IP地址
            copy_ip_address(info->ip_addr, WiFi.localIP());
            copy_ip_address(info->gateway, WiFi.gatewayIP());
            copy_ip_address(info->netmask, WiFi.subnetMask());
            copy_ip_address(info->dns1, WiFi.dnsIP(0));
            copy_ip_address(info->dns2, WiFi.dnsIP(1));
            
            // 获取MAC地址
            uint8_t mac[6];
            WiFi.macAddress(mac);
            memcpy(info->mac, mac, 6);
            
            // 获取SSID
            String ssid = WiFi.SSID();
            strncpy(info->ssid, ssid.c_str(), sizeof(info->ssid) - 1);
            
            // 获取信道
            info->channel = WiFi.channel();
            
            // 获取信号强度
            info->rssi = WiFi.RSSI();
            
            return OK;
        }
    } 
    else if (mode == WLAN_MODE_AP) {
        // AP 模式
        if (WiFi.softAPgetStationNum() > 0) {
            info->connected = 1;
            
            // 修复AP模式下的IP地址拷贝
            copy_ip_address(info->ip_addr, WiFi.softAPIP());
            copy_ip_address(info->gateway, WiFi.softAPIP());  // AP模式下网关就是自己
            
            IPAddress netmask(255, 255, 255, 0);
            copy_ip_address(info->netmask, netmask);
            
            // 获取MAC地址
            uint8_t mac[6];
            WiFi.softAPmacAddress(mac);
            memcpy(info->mac, mac, 6);
            
            // 获取SSID
            String ssid = WiFi.softAPSSID();
            strncpy(info->ssid, ssid.c_str(), sizeof(info->ssid) - 1);
            
            // 获取信道
            info->channel = WiFi.channel();
            
            return OK;
        }
    }
    
    info->connected = 0;
    return OK;
}

// 更新WLAN操作结构体
static const WlanOps esp32_wlan_ops = {
    .set_mode = esp32_set_mode,
    .get_mode = esp32_get_mode,
    .scan = esp32_scan,
    .connect = esp32_connect,
    .startap = esp32_startap,
    .disconnect = esp32_disconnect,
    .get_net_info = esp32_get_net_info,
    .read_config = esp32_read_config,
    .save_config = esp32_save_config,
    .get_mac = esp32_get_mac
};

// Initialize ESP32 WLAN driver
void Esp32WlanInit(void)
{
    // Initialize WiFi
    WiFi.persistent(false);
    WiFi.mode(WIFI_OFF);
    
    g_e2p = (E2prom*)ios_open("e2p", 0, 0);

    // Create WLAN device
    WlanCreate(&_wlan, "wlan0", &esp32_wlan_ops);
    _wlan.Protected.config = &_config;
    
    SYSLOG(LOG_INFO, "ESP32 WLAN driver initialized");
}

// Example usage commands
#include "debug.h"
#include "console.h"

static enConsoleResult wlan_scan(int argc, char** argv)
{
    // Set WLAN mode to STATION before scanning
    if(WlanSetMode(&_wlan, WLAN_MODE_STATION) != OK) {
        KPrint("Failed to set station mode for scanning\r\n");
        return CE_RUNTIME_ERROR;
    }

    WlanScanInfo scan_info[20];
    _config.scan_results = scan_info;
    _config.scan_count = 20;
    
    int count = WlanScan(&_wlan, scan_info, 20);
    if(count > 0) {
        KPrintf("\r\nFound %d networks:\r\n", count);
        for(int i = 0; i < count; i++) {
            KPrintf("%d: SSID: %s, BSSID: %02X:%02X:%02X:%02X:%02X:%02X, Channel: %d\r\n",
                i + 1,
                scan_info[i].ssid,
                scan_info[i].bssid[0], scan_info[i].bssid[1], scan_info[i].bssid[2],
                scan_info[i].bssid[3], scan_info[i].bssid[4], scan_info[i].bssid[5],
                scan_info[i].channel);
        }
    } else {
        KPrint("Scan failed\r\n");
    }
    return CE_OK;
}

static enConsoleResult wlan_connect(int argc, char** argv)
{
    if(argc < 3 || argc > 4) {
        KPrint("Usage: wlan_connect <ssid> <password> [timeout_sec]\r\n");
        KPrint("       timeout_sec: connection timeout in seconds (default: 30)\r\n");
        return CE_PARAM_ERROR;
    }

    uint32_t timeout_sec = 30; // 默认30秒超时
    if(argc == 4) {
        timeout_sec = atoi(argv[3]);
        if(timeout_sec == 0) {
            timeout_sec = 30; // 如果解析失败或输入0，使用默认值
        }
    }

    if(WlanSetMode(&_wlan, WLAN_MODE_STATION) != OK) {
        KPrint("Failed to set station mode\r\n");
        return CE_RUNTIME_ERROR;
    }

    if(WlanConnect(&_wlan, argv[1], argv[2], timeout_sec) == OK) {
        KPrint("Connected to WiFi\r\n");
        KPrintf("IP: %s\r\n", WiFi.localIP().toString().c_str());
        return CE_OK;
    }

    KPrint("Failed to connect\r\n");
    return CE_RUNTIME_ERROR;
}

static enConsoleResult wlan_ap(int argc, char** argv)
{
    if(argc != 3) {
        KPrint("Usage: wlan_ap <ssid> <password>\r\n");
        return CE_PARAM_ERROR;
    }

    if(WlanSetMode(&_wlan, WLAN_MODE_AP) != OK) {
        KPrint("Failed to set AP mode\r\n");
        return CE_RUNTIME_ERROR;
    }

    if(WlanStartAP(&_wlan, argv[1], argv[2]) == OK) {
        KPrint("AP started\r\n");
        KPrintf("AP IP: %s\r\n", WiFi.softAPIP().toString().c_str());
        return CE_OK;
    }

    KPrint("Failed to start AP\r\n");
    return CE_RUNTIME_ERROR;
}

static enConsoleResult wlan_config(int argc, char** argv) {
    if (argc < 2) {
        KPrint("Usage: wlan_config <save|clear|show>\r\n");
        return CE_PARAM_ERROR;
    }

    char ssid[33] = {0};
    char password[65] = {0};

    if (strcmp(argv[1], "save") == 0) {
        // 保存网络配置
        if (argc != 4) {
            KPrint("Usage: wlan_config save <ssid> <password>\r\n");
            return CE_PARAM_ERROR;
        }
        const char* current_ssid = argv[2];
        const char* current_pass = argv[3];
        if (esp32_save_config(&_wlan, current_ssid, current_pass) == OK) {
            KPrint("Current network configuration saved\r\n");
            return CE_OK;
        }
    }
    else if (strcmp(argv[1], "clear") == 0) {
        // 清除配置（保存空配置）
        if (esp32_save_config(&_wlan, "", "") == OK) {
            KPrint("Network configuration cleared\r\n");
            return CE_OK;
        }
    }
    else if (strcmp(argv[1], "show") == 0) {
        // 显示当前保存的配置
        if (esp32_read_config(&_wlan, ssid, password) == OK) {
            KPrint("Saved configuration:\r\n");
            KPrintf("SSID: %s\r\n", ssid);
            if (strlen(password) > 0) {
                KPrint("Password: ********\r\n");
            } else {
                KPrint("Password: <none>\r\n");
            }
            return CE_OK;
        }
    }

    KPrint("Operation failed\r\n");
    return CE_RUNTIME_ERROR;
}

static enConsoleResult ping(int argc, char** argv)
{
    if (argc != 2) {
        KPrint("Usage: ping <host>\r\n");
        return CE_PARAM_ERROR;
    }

    // 检查WiFi是否连接
    if (WiFi.status() != WL_CONNECTED) {
        KPrint("Error: WiFi not connected\r\n");
        return CE_RUNTIME_ERROR;
    }

    const char* target = argv[1];
    bool is_ip = false;
    IPAddress target_ip;

    // 检查是否是IP地址
    if (target_ip.fromString(target)) {
        is_ip = true;
    }

    KPrintf("\nPING %s\n", target);

    int success = 0;
    int failures = 0;
    int total_time = 0;
    const int count = 4;  // ping次数

    for (int i = 0; i < count; i++) {
        int avg_time_ms = 0;
        bool ret;
        
        if (is_ip) {
            ret = Ping.ping(target_ip, 1);
            avg_time_ms = Ping.averageTime();
        } else {
            ret = Ping.ping(target, 1);
            avg_time_ms = Ping.averageTime();
        }

        if (ret) {
            KPrintf("64 bytes from %s: icmp_seq=%d time=%d ms\n", 
                    target, i + 1, avg_time_ms);
            success++;
            total_time += avg_time_ms;
        } else {
            KPrintf("From %s: icmp_seq=%d timeout\n", 
                    target, i + 1);
            failures++;
        }

        // 等待1秒再发送下一个包
        if (i < count - 1) {
            atomTimerDelay(100);
        }
    }

    // 打印统计信息
    KPrintf("\n--- %s ping statistics ---\n", target);
    KPrintf("%d packets transmitted, %d received, %d%% packet loss", 
            count, success, (failures * 100) / count);
    
    if (success > 0) {
        KPrintf(", avg time %d ms", total_time / success);
    }
    KPrintf("\n");

    return CE_OK;
}

static enConsoleResult wlan_disconnect(int argc, char** argv)
{
    if (WlanDisconnect(&_wlan) == OK) {
        KPrint("Disconnected from WiFi\r\n");
        return CE_OK;
    }

    KPrint("Failed to disconnect\r\n");
    return CE_RUNTIME_ERROR;
}

static enConsoleResult netstat(int argc, char** argv)
{
    WlanNetInfo info;
    if (WlanGetNetInfo(&_wlan, &info) != OK) {
        KPrint("Failed to get network information\r\n");
        return CE_RUNTIME_ERROR;
    }

    KPrint("\r\nNetwork Status:\r\n");
    KPrintf("Connection: %s\r\n", info.connected ? "Connected" : "Disconnected");
    
    if (info.connected) {
        KPrintf("IP Address: %d.%d.%d.%d\r\n", 
                info.ip_addr[0], info.ip_addr[1], info.ip_addr[2], info.ip_addr[3]);
        KPrintf("Gateway: %d.%d.%d.%d\r\n",
                info.gateway[0], info.gateway[1], info.gateway[2], info.gateway[3]);
        KPrintf("Netmask: %d.%d.%d.%d\r\n",
                info.netmask[0], info.netmask[1], info.netmask[2], info.netmask[3]);
        KPrintf("DNS1: %d.%d.%d.%d\r\n",
                info.dns1[0], info.dns1[1], info.dns1[2], info.dns1[3]);
        KPrintf("DNS2: %d.%d.%d.%d\r\n",
                info.dns2[0], info.dns2[1], info.dns2[2], info.dns2[3]);
        KPrintf("MAC: %02X:%02X:%02X:%02X:%02X:%02X\r\n",
                info.mac[0], info.mac[1], info.mac[2],
                info.mac[3], info.mac[4], info.mac[5]);
        KPrintf("SSID: %s\r\n", info.ssid);
        KPrintf("Channel: %d\r\n", info.channel);
        KPrintf("RSSI: %d dBm\r\n", info.rssi);
    }
    
    return CE_OK;
}

BGEIN_DEF_CMD(_wlanCmds)
    CMD_DEF(wlan_scan, "scan for wifi networks")
    CMD_DEF(wlan_connect, "connect to wifi network")
    CMD_DEF(wlan_ap, "start wifi access point")
    CMD_DEF(wlan_config, "manage wifi configuration")
    CMD_DEF(netstat, "show WiFi AP info and connected clients")
    CMD_DEF(ping, "ping a host")
    CMD_DEF(wlan_disconnect, "disconnect from wifi network")
END_DEF_CMD()


void Esp32WlanCreate()
{
    Esp32WlanInit();

    CONSOLE_REG_CMD(_wlanCmds);
}
